import os
import numpy as np
import pandas as pd
import glob
from pathlib import Path
import json


def json_filter(json_path, dst_path):
    with open(json_path, 'r') as f:
        content = f.read()
        data = json.loads(content)
    
    data['imageData'] = 'null'
    
    new_data_shapes= []
    for d in data['shapes']:
        d['label'] = d['label'].replace(' ', '')
        if len(d['points']) < 3:
            continue
        if len(d['label']) == 6 and d['label'][-1] == '1':  # 去掉前程修补
            continue
        if len(d['label']) == 6:  # 合并2 & 3至主code
            d['label'] = d['label'][:-1]
        new_data_shapes.append(d)
    data['shapes'] = new_data_shapes

    json_name = str(Path(json_path).name)
    code = str(Path(json_path).parent.name)
    json_dst_path = os.path.join(dst_path, code, json_name)
    os.makedirs(Path(json_dst_path).parent, exist_ok=True)
    with open(json_dst_path, 'w') as f:
        b = json.dumps(data, indent=4)
        f.write(b)


def main():
    src_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_img_cut_pad'
    dst_path = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/filtered_json'

    df = pd.DataFrame({'json':glob.glob(os.path.join(src_path, '*/*.json'))})

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False
    
    if is_pandarallel:
        df['json'].parallel_apply(lambda x: json_filter(x, dst_path))
    else:
        df['json'].apply(lambda x: json_filter(x, dst_path))



if __name__=='__main__':
    main()
